"""
python p2_solution_acoant.py
"""

import json
import os
import numpy as np
import networkx as nx
from acopy import Colony, Solver

# 路径设置
warehouse_path = '../fujian/fujian3/origin_data/warehouse.json'
inventory_path = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
sales_path = '../fujian/fujian3/data_from_p1/all_average_sales.json'
output_path = '../fujian/fujian3/ant/'

# 数据加载
with open(warehouse_path, 'r') as f:
    warehouses = json.load(f)
with open(inventory_path, 'r') as f:
    inventory_data = json.load(f)
with open(sales_path, 'r') as f:
    sales_data = json.load(f)

# 数据转换
warehouse_dict = {w['warehouse_id']: w for w in warehouses}
inventory_dict = {i['category_id']: i['average_inventory'] for i in inventory_data}
sales_dict = {s['category_id']: s['average_sales'] for s in sales_data}

# 创建图对象
graph = nx.DiGraph()

# 添加节点到图中
for warehouse_id in warehouse_dict.keys():
    graph.add_node(warehouse_id)

# 配置蚁群算法参数
colony = Colony(alpha=1.0, beta=2.0, graph=graph)  # 传入图对象
solver = Solver(rho=0.5, q=1.0)

# 定义目标函数
def objective_function(assignment):
    total_cost = 0
    warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouse_dict.keys()}
    
    # 计算成本和利用率
    for category, warehouse in assignment.items():
        inventory = inventory_dict[category]
        sales = sales_dict[category]
        warehouse_info = warehouse_dict[warehouse]
        
        # 更新仓库数据
        warehouse_utilization[warehouse]['inventory'] += inventory
        warehouse_utilization[warehouse]['sales'] += sales

    # 检查约束条件
    for warehouse_id, usage in warehouse_utilization.items():
        if (usage['inventory'] >= warehouse_dict[warehouse_id]['max_inventory'] or 
            usage['sales'] >= warehouse_dict[warehouse_id]['max_sales']):
            return float('inf')  # 违反约束条件

    # 满足约束条件时计算总成本
    total_cost = sum(warehouse_dict[w]['daily_cost'] for w in warehouse_utilization if warehouse_utilization[w]['inventory'] > 0)
    return total_cost

# 定义路径并求解
def solve_warehouse_assignment():
    categories = list(inventory_dict.keys())
    best_solution = solver.solve(graph, objective_function)  # 传递图对象
    return best_solution

# 生成解和保存结果
solution = solve_warehouse_assignment()
assignment = {category: solution.path[category] for category in categories}

# 计算总成本和利用率
total_cost = objective_function(assignment)
utilization_rates = []
for warehouse_id, usage in warehouse_utilization.items():
    utilization_rates.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_inventory': usage['inventory'] / warehouse_dict[warehouse_id]['max_inventory'],
        'utilization_rate_of_sales': usage['sales'] / warehouse_dict[warehouse_id]['max_sales']
    })

# 保存结果
os.makedirs(output_path, exist_ok=True)
with open(os.path.join(output_path, 'all_cost.txt'), 'w') as f:
    f.write(str(total_cost))

with open(os.path.join(output_path, 'all_utilization_rate_inventory.json'), 'w') as f:
    json.dump(utilization_rates, f)

with open(os.path.join(output_path, 'assignment.json'), 'w') as f:
    json.dump([{'category_id': k, 'warehouse_id': v} for k, v in assignment.items()], f)
